import sqlite3
from PyQt5.QtWidgets import QMainWindow, QMessageBox, QMenu, QAction, QTableWidgetItem
from PyQt5.QtCore import Qt, QPoint
from AggregateWin import Ui_MainWindow as AggregateWinUI
from PyQt5.QtGui import QIcon
from datetime import datetime
from DetailsWin import Ui_MainWindow as DetailsWinUI
import re
import urllib.parse

# 全局窗口管理器，用于存储和复用窗口实例
_window_instances = {}

class AggregateWindowManager:
    """聚合窗口管理器，用于创建和复用窗口实例"""
    
    def __init__(self):
        self.windows = {}
    
    def get_window(self, window_id):
        """获取指定ID的窗口实例，如果不存在则创建"""
        if window_id not in self.windows or self.windows[window_id] is None:
            # 创建新窗口
            aggregate_window = QMainWindow()
            ui = AggregateWinUI()
            ui.setupUi(aggregate_window)
            
            # 存储窗口实例
            self.windows[window_id] = {
                'window': aggregate_window,
                'ui': ui
            }
        
        return self.windows[window_id]
    
    def update_window(self, window_id, title, data, columns, column_headers):
        """更新窗口内容"""
        try:
            # 获取或创建窗口
            window_data = self.get_window(window_id)
            window = window_data['window']
            ui = window_data['ui']
            
            # 设置窗口标题
            window.setWindowTitle(title)
            
            # 设置表格列数和表头
            ui.tableWidget.setColumnCount(len(column_headers))
            ui.tableWidget.setHorizontalHeaderLabels(column_headers)
            
            # 设置表格行数
            ui.tableWidget.setRowCount(len(data))
            
            # 保存列信息和数据，供右键菜单使用
            ui._columns = columns
            ui._data = data
            ui._window_id = window_id
            
            # 填充表格数据
            for row_idx, row_data in enumerate(data):
                for col_idx, cell_data in enumerate(row_data):
                    # 处理特殊数据格式
                    display_data = self._format_cell_data(cell_data, col_idx, columns)
                    
                    # 创建单元格项
                    item = QTableWidgetItem(display_data)
                    
                    # 设置对齐方式
                    if columns[col_idx] == 'url':
                        item.setTextAlignment(Qt.AlignLeft | Qt.AlignVCenter)
                    else:
                        item.setTextAlignment(Qt.AlignCenter)
                    
                    ui.tableWidget.setItem(row_idx, col_idx, item)
            
            # 自适应列宽
            from PyQt5 import QtWidgets
            ui.tableWidget.horizontalHeader().setSectionResizeMode(QtWidgets.QHeaderView.Stretch)
            ui.tableWidget.verticalHeader().setVisible(False)
            window.setWindowIcon(QIcon("logo.jpg"))
            
            # 设置右键菜单
            self._setup_context_menu(ui)
            
            # 显示窗口
            window.show()
            window.raise_()  # 窗口置顶
            
            return window
            
        except Exception as e:
            print(f"更新窗口时出错: {e}")
            import traceback
            traceback.print_exc()
            return None
    
    def _setup_context_menu(self, ui):
        """设置表格右键菜单"""
        # 创建右键菜单
        if hasattr(ui, 'context_menu'):
            # 断开之前的信号连接
            try:
                ui.tableWidget.customContextMenuRequested.disconnect()
            except TypeError:
                pass  # 没有连接时忽略错误
            
            # 清除菜单项
            ui.context_menu.clear()
        else:
            # 创建新菜单
            ui.context_menu = QMenu(ui.tableWidget)
        
        # 创建菜单项
        ui.action_ip_request_list = QAction("请求列表", ui.tableWidget)
        ui.action_ip_red = QAction("IP标红", ui.tableWidget)
        ui.action_ip_green = QAction("IP标绿", ui.tableWidget)
        ui.action_aggregate_url = QAction("URL聚合", ui.tableWidget)
        ui.action_aggregate_time = QAction("时间聚合", ui.tableWidget)
        ui.action_all_info = QAction("详细信息", ui.tableWidget)
        
        # 添加菜单项到菜单
        ui.context_menu.addAction(ui.action_ip_request_list)
        ui.context_menu.addAction(ui.action_ip_red)
        ui.context_menu.addAction(ui.action_ip_green)
        ui.context_menu.addSeparator()
        ui.context_menu.addAction(ui.action_aggregate_url)
        ui.context_menu.addAction(ui.action_aggregate_time)
        ui.context_menu.addSeparator()
        ui.context_menu.addAction(ui.action_all_info)
        
        # 连接信号槽
        ui.action_ip_request_list.triggered.connect(lambda: self._on_action_triggered(ui, "ip_request_list"))
        ui.action_ip_red.triggered.connect(lambda: self._on_action_triggered(ui, "ip_red"))
        ui.action_ip_green.triggered.connect(lambda: self._on_action_triggered(ui, "ip_green"))
        ui.action_aggregate_url.triggered.connect(lambda: self._on_action_triggered(ui, "aggregate_url"))
        ui.action_aggregate_time.triggered.connect(lambda: self._on_action_triggered(ui, "aggregate_time"))
        ui.action_all_info.triggered.connect(lambda: self._on_action_triggered(ui, "all_info"))
        
        # 设置表格的右键菜单策略
        ui.tableWidget.setContextMenuPolicy(Qt.CustomContextMenu)
        ui.tableWidget.customContextMenuRequested.connect(lambda point: self._show_context_menu(ui, point))
    
    def _show_context_menu(self, ui, point):
        """显示右键菜单"""
        # 检查是否有选中的单元格
        selected_items = ui.tableWidget.selectedItems()
        if selected_items:
            # 有选中条目，显示菜单
            ui.context_menu.exec_(ui.tableWidget.mapToGlobal(point))
    
    def _on_action_triggered(self, ui, action_type):
        """处理菜单项触发事件"""
        # 获取选中的行
        selected_row = ui.tableWidget.currentRow()
        if selected_row >= 0:
            # 尝试从原始数据中获取IP、URL和时间戳
            # 假设这些字段在原始数据中的位置
            ip = None
            url = None
            timestamp = None
            
            # 查找IP、URL和timestamp在列中的索引
            for col_idx, col_name in enumerate(ui._columns):
                if col_name == 'ip' and selected_row < len(ui._data) and col_idx < len(ui._data[selected_row]):
                    ip = ui._data[selected_row][col_idx]
                elif col_name == 'url' and selected_row < len(ui._data) and col_idx < len(ui._data[selected_row]):
                    url = ui._data[selected_row][col_idx]
                elif col_name == '时间' and selected_row < len(ui._data) and col_idx < len(ui._data[selected_row]):
                    timestamp = ui._data[selected_row][col_idx]
            
            # 如果找不到，尝试从表格中获取
            if not ip and ui.tableWidget.item(selected_row, 0):
                ip = ui.tableWidget.item(selected_row, 0).text()
            if not url and ui.tableWidget.columnCount() > 2 and ui.tableWidget.item(selected_row, 2):
                url = ui.tableWidget.item(selected_row, 2).text()
            if not timestamp and ui.tableWidget.columnCount() > 5 and ui.tableWidget.item(selected_row, 5):
                timestamp = ui.tableWidget.item(selected_row, 5).text()
            
            # 根据操作类型调用相应的函数
            if action_type == "ip_request_list" and ip:
                # 保存窗口实例，避免被垃圾回收
                if not hasattr(self, '_window_instances'):
                    self._window_instances = {}
                self._window_instances[f"request_list_{ip}"] = request_list(ip)
            elif action_type == "ip_red" and ip:
                mark_ip_red(ip)
            elif action_type == "ip_green" and ip:
                mark_ip_green(ip)
            elif action_type == "aggregate_url" and url:
                # 保存窗口实例
                if not hasattr(self, '_window_instances'):
                    self._window_instances = {}
                self._window_instances[f"aggregate_url_{url[:30]}"] = aggregate_by_url(url)
            elif action_type == "aggregate_time" and timestamp:
                # 保存窗口实例
                if not hasattr(self, '_window_instances'):
                    self._window_instances = {}
                self._window_instances[f"aggregate_time_{timestamp[:20]}"] = aggregate_by_time(timestamp)
            elif action_type == "all_info":
                # 尝试从数据库中查找匹配的记录ID
                if ip and url:
                    conn = sqlite3.connect('nginx_logs.db')
                    cursor = conn.cursor()
                    cursor.execute("SELECT id FROM nginx_logs WHERE ip = ? AND string = ? LIMIT 1", (ip, url))
                    row = cursor.fetchone()
                    conn.close()
                    if row:
                        id = row[0]
                        # 保存窗口实例
                        if not hasattr(self, '_window_instances'):
                            self._window_instances = {}
                        self._window_instances[f"all_info_{id}"] = all_info(id)
    
    def _format_cell_data(self, data, col_idx, columns):
        """格式化单元格数据"""
        if data is None:
            return ""
            
        # 处理百分比列
        if columns[col_idx] in ['point', 'max_point', 'avg_point', 'avg_risk']:
            try:
                return f"{float(data) * 100:.2f}%"
            except (ValueError, TypeError):
                return str(data)
        # 处理时间戳列
        elif columns[col_idx] in ['timestamp', 'first_timestamp', 'last_timestamp', 'first_request', 'last_request', 'hour']:
            try:
                import datetime
                if isinstance(data, str) and ':' in data:
                    # 已经格式化的时间字符串
                    return data
                timestamp = float(data)
                dt = datetime.datetime.fromtimestamp(timestamp)
                return dt.strftime('%Y-%m-%d %H:%M:%S')
            except (ValueError, TypeError):
                return str(data)
        else:
            return str(data)

# 创建全局窗口管理器实例
window_manager = AggregateWindowManager()

def request_list(ip):
    """获取指定IP的请求列表并在AggregateWin窗口中显示"""
    try:
        # 连接数据库查询该IP的所有请求记录
        conn = sqlite3.connect('nginx_logs.db')
        cursor = conn.cursor()
        
        # 查询指定IP的所有请求记录
        cursor.execute("""
            SELECT ip, address, string, point, timestamp 
            FROM nginx_logs 
            WHERE ip = ? 
            ORDER BY timestamp DESC
        """, (ip,))
        rows = cursor.fetchall()
        conn.close()
        
        # 定义列信息
        columns = ['ip', 'address', 'url', 'point', 'timestamp']
        column_headers = ['IP', '归属', 'URL', '攻击概率(%)', '时间']
        
        # 使用窗口管理器更新窗口
        window_id = f"request_list_{ip}"
        title = f"IP: {ip} 的请求列表"
        
        return window_manager.update_window(window_id, title, rows, columns, column_headers)
        
    except Exception as e:
        print(f"显示请求列表时出错: {e}")
        import traceback
        traceback.print_exc()
        return None

def mark_ip_red(ip):
    """将指定IP标红 - 更新数据库中该IP的所有条目的lable列为1"""
    try:
        # 连接数据库
        conn = sqlite3.connect('nginx_logs.db')
        cursor = conn.cursor()
        
        # 更新指定IP的所有条目的lable列为1
        cursor.execute("UPDATE nginx_logs SET lable = 1 WHERE ip = ?", (ip,))
        
        # 获取受影响的行数
        affected_rows = cursor.rowcount
        
        # 提交事务
        conn.commit()
        conn.close()
        
        print(f"已将IP {ip} 的 {affected_rows} 条记录标红(lable=1)")
        return affected_rows
        
    except Exception as e:
        print(f"标红IP {ip} 时出错: {e}")
        import traceback
        traceback.print_exc()
        return 0
    

def mark_ip_green(ip):
    """将指定IP标绿 - 更新数据库中该IP的所有条目的lable列为0"""
    try:
        # 连接数据库
        conn = sqlite3.connect('nginx_logs.db')
        cursor = conn.cursor()
        
        # 更新指定IP的所有条目的lable列为1
        cursor.execute("UPDATE nginx_logs SET lable = 0 WHERE ip = ?", (ip,))
        
        # 获取受影响的行数
        affected_rows = cursor.rowcount
        
        # 提交事务
        conn.commit()
        conn.close()
        
        print(f"已将IP {ip} 的 {affected_rows} 条记录标绿(lable=0)")
        return affected_rows
        
    except Exception as e:
        print(f"标绿IP {ip} 时出错: {e}")
        import traceback
        traceback.print_exc()
        return 0
        

def aggregate_by_url(url):
    """显示所有匹配指定URL的记录"""
    try:
        # 连接数据库
        conn = sqlite3.connect('nginx_logs.db')
        cursor = conn.cursor()
        
        # 查询所有匹配该URL的记录，按时间降序排序
        cursor.execute("""
            SELECT 
                ip, 
                address,
                string as url, 
                point,
                timestamp 
            FROM nginx_logs 
            WHERE string = ? 
            ORDER BY timestamp DESC
        """, (url,))
        rows = cursor.fetchall()
        conn.close()
        
        # 定义列信息
        columns = ['ip', 'address', 'url', 'point', 'timestamp']
        column_headers = ['IP', '归属', 'URL', '攻击概率(%)', '请求时间']
        
        # 使用窗口管理器更新窗口
        window_id = f"aggregate_url_{url[:30]}"
        title = f"URL: {url[:50]}... 的请求记录"
        
        return window_manager.update_window(window_id, title, rows, columns, column_headers)
        
    except Exception as e:
        print(f"URL记录查询时出错: {e}")
        import traceback
        traceback.print_exc()
        return None


def aggregate_by_time(timerange):
    """显示指定时间附近4小时范围内的所有请求记录"""
    print(f"aggregate_by_time: {timerange}")
    try:
        # 解析时间范围
        import datetime
        timestamp = None
        try:
            # 首先尝试直接转换为浮点数（如果已经是时间戳格式）
            try:
                base_time = datetime.datetime.fromtimestamp(timestamp)
            except (ValueError, TypeError):
                # 如果是人类可读的时间字符串格式，进行解析
                # 假设格式为 "2025-08-09 10:14:26"
                try:
                    base_time = datetime.datetime.strptime(timerange.strip(), '%Y-%m-%d %H:%M:%S')
                    timestamp = base_time.timestamp()
                except ValueError:
                    # 尝试分割并解析第一个部分
                    time_part = timerange.split()[0] if isinstance(timerange, str) else str(timerange)
                    try:
                        base_time = datetime.datetime.strptime(time_part, '%Y-%m-%d')
                    except ValueError:
                        base_time = datetime.datetime.strptime(time_part, '%Y-%m-%d %H:%M:%S')
                    timestamp = base_time.timestamp()
                    
            # 计算4小时时间范围（中心时间±2小时）
            start_time = base_time - datetime.timedelta(hours=2)
            end_time = base_time + datetime.timedelta(hours=2)
            time_str = base_time.strftime('%Y-%m-%d %H:%M:%S')
            print(f"解析成功，中心时间: {time_str}，时间戳: {timestamp}")
            print(f"查询范围: {start_time} 到 {end_time}")
            
        except Exception as e:
            print(f"时间解析失败: {e}")
            # 如果无法解析时间，使用当前时间附近4小时
            base_time = datetime.datetime.now()
            start_time = base_time - datetime.timedelta(hours=2)
            end_time = base_time + datetime.timedelta(hours=2)
            time_str = "当前时间"
            print(f"使用当前时间: {time_str}")
        
        # 连接数据库
        conn = sqlite3.connect('nginx_logs.db')
        cursor = conn.cursor()
        
        # 查询指定时间范围内的所有记录，按时间戳降序排序
        cursor.execute("""
            SELECT 
                ip, 
                address,
                string as url, 
                point,
                timestamp 
            FROM nginx_logs 
            WHERE timestamp >= ? AND timestamp <= ? 
            ORDER BY timestamp DESC
        """, (start_time.timestamp(), end_time.timestamp()))
        rows = cursor.fetchall()
        conn.close()
        
        # 定义列信息
        columns = ['ip', 'address', 'url', 'point', 'timestamp']
        column_headers = ['IP', '地址', 'URL', '攻击概率(%)', '请求时间']   
        
        # 使用窗口管理器更新窗口
        window_id = f"aggregate_time_{time_str[:10]}"
        title = f"{time_str} 附近4小时请求记录"
        
        return window_manager.update_window(window_id, title, rows, columns, column_headers)
        
    except Exception as e:
        print(f"时间范围内记录查询时出错: {e}")
        import traceback
        traceback.print_exc()
        return None

def parse_log_string(log_string):
    """解析日志字符串，提取URL、状态码、响应长度等信息"""
    # 尝试解析日志格式，提取URL、方法、状态码等信息
    result = {
        'method': '',
        'url': '',
        'protocol': '',
        'status_code': '',
        'response_size': '',
        'referer': '',
        'user_agent': ''
    }
    
    # 尝试匹配HTTP请求行
    request_pattern = r'^(\S+)\s+(.*?)\s+(HTTP/\d\.\d)'
    match = re.match(request_pattern, log_string)
    if match:
        result['method'] = match.group(1)
        result['url'] = match.group(2)
        result['protocol'] = match.group(3)
    
    return result


def url_decode(ui):
        """
        对plainTextEdit_url中的内容进行URL解码
        """
        
        try:
            # 执行URL解码
            decoded_url = urllib.parse.unquote(ui.plainTextEdit_url.toPlainText())
            
            # 如果解码后内容发生变化，更新到文本框
            if decoded_url != ui.plainTextEdit_url.toPlainText():
                ui.plainTextEdit_url.setPlainText(decoded_url)
            else:
                # 如果没有变化，可能是多层编码，尝试再次解码
                double_decoded_url = urllib.parse.unquote(decoded_url)
                if double_decoded_url != decoded_url:
                    ui.plainTextEdit_url.setPlainText(double_decoded_url)
        except Exception as e:
            # 显示错误信息
            print(f"URL解码错误: {e}")
            QMessageBox.warning(
                None, 
                "解码错误", 
                f"URL解码过程中发生错误: {str(e)}"
            )
            ui.plainTextEdit_url.setPlainText(ui.plainTextEdit_url.toPlainText())
def all_info(id):
    """查看详情"""
    try:
        # 连接数据库
        conn = sqlite3.connect('nginx_logs.db')
        cursor = conn.cursor()
        
        # 查询指定ID的记录
        cursor.execute("SELECT ip, string, point, timestamp, original_log, address, regular_result FROM nginx_logs WHERE id = ?", (id,))
        row = cursor.fetchone()
        conn.close()
        
        if not row:
            print(f"未找到ID为 {id} 的记录")
            return None
        
        # 解包数据
        ip, log_string, point, timestamp, original_log, address, regular_result = row
        
        # 解析日志字符串
        parsed_info = parse_log_string(log_string)
        
        # 格式化时间
        try:
            dt = datetime.fromtimestamp(int(timestamp))
            formatted_time = dt.strftime('%Y-%m-%d %H:%M:%S')
        except:
            formatted_time = '未知'
        
        # 创建新窗口实例
        details_window = QMainWindow()
        ui = DetailsWinUI()
        ui.setupUi(details_window)
        from PyQt5.QtGui import QIcon
        details_window.setWindowIcon(QIcon("logo.jpg"))
        ui.pushButton.clicked.connect(lambda: url_decode(ui))
        # 填充数据
        ui.lineEdit_ip.setText(ip)
        ui.lineEdit_address.setText(address or '未知')
        ui.lineEdit_formattime.setText(formatted_time)
        ui.lineEdit_time.setText(str(timestamp) if timestamp else '未知')
        ui.plainTextEdit_url.setPlainText(parsed_info['url'] or log_string)
        
        # 尝试从原始日志中提取更多信息
        status_match = re.search(r'\s+(\d{3})\s+(\d+|"-"|"")', original_log)
        if status_match:
            ui.lineEdit_statuscode.setText(status_match.group(1) if len(status_match.groups()) > 0 else '')
            ui.lineEdit_length.setText(status_match.group(2) if len(status_match.groups()) > 1 else '')
        
        # 尝试提取Referer和User-Agent
        referer_match = re.search(r'"Referer: ([^"]*)"', original_log)
        if referer_match:
            ui.plainTextEdit_refere.setPlainText(referer_match.group(1))
        
        # 尝试找到User-Agent部分
        parts = original_log.split('"')
        if len(parts) >= 7:
            ui.plainTextEdit_ua.setPlainText(parts[5])
        
        # 设置原始日志
        ui.plainTextEdit_log.setPlainText(original_log)
        
        # 设置AI预测结果
        try:
            probability = float(point)
            ui.lineEdit_ai_result.setText(f"{probability:.2%}")
        except:
            ui.lineEdit_ai_result.setText("未知")
        
        # 设置正则别结果
        ui.lineEdit_other_result.setText("未判断" if regular_result is None else ("攻击" if regular_result == '1' else "正常"))
        
        # 显示窗口
        details_window.setWindowTitle(f"日志详情 - ID: {id}")
        details_window.show()
        
        print(f"成功显示ID为 {id} 的记录详情")
        return details_window
        
    except Exception as e:
        print(f"显示详情时出错: {e}")
        import traceback
        traceback.print_exc()
        return None